"""
1. 利用langchain 对文本进行切片
2. 将切片的内容添加到向量数据库中,向量数据库使用milvus
3. 进行相似性检索
4. 返回相似性检索的结果
"""

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_milvus import MilvusVectorStore
from langchain_openai import OpenAIEmbeddings
import os
import time
from pymilvus import connections, utility

class TextSearchEngine:
    def __init__(self, connection_args=None):
        # 默认连接参数
        self.connection_args = connection_args or {
            "host": "localhost",
            "port": "19530"
        }
        
        # 尝试连接Milvus
        try:
            connections.connect(**self.connection_args)
            print("成功连接到Milvus服务器")
        except Exception as e:
            print(f"连接Milvus服务器失败: {str(e)}")
            print("请确保Milvus服务器已启动且可访问")
            raise
        
        # 初始化文本分割器，使用更小的chunk_size
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=10,  # 减小chunk大小
            chunk_overlap=2  # 减小重叠大小
        )
        
        # 初始化embedding模型
        api_key = os.getenv("SILICONFLOW_API_TOKEN")
        if not api_key:
            raise ValueError("请设置环境变量 SILICONFLOW_API_TOKEN")
            
        self.embeddings = OpenAIEmbeddings(
            model="BAAI/bge-large-zh-v1.5",
            api_key=api_key,
            base_url="https://api.siliconflow.cn/v1",
        )
        
    def process_and_store(self, text, collection_name="text_collection"):
        """处理文本并存储到Milvus中"""
        try:
            # 切分文本
            texts = self.text_splitter.split_text(text)
            print(f"文本被切分为 {len(texts)} 个片段")
            
            # 创建Milvus向量存储
            vector_store = MilvusVectorStore.from_texts(
                texts=texts,
                embedding=self.embeddings,
                collection_name=collection_name,
                connection_args=self.connection_args
            )
            
            return vector_store
        except Exception as e:
            print(f"处理文本时发生错误: {str(e)}")
            raise
    
    def search_similar(self, query, collection_name="text_collection", top_k=3):
        """搜索相似内容"""
        try:
            # 连接已存在的collection
            vector_store = MilvusVectorStore(
                embedding_function=self.embeddings,
                collection_name=collection_name,
                connection_args=self.connection_args
            )
            
            # 执行相似性搜索
            results = vector_store.similarity_search_with_score(query, k=top_k)
            return results
        except Exception as e:
            print(f"搜索时发生错误: {str(e)}")
            raise

def main():
    # 使用示例
    search_engine = TextSearchEngine()
    
    # 示例文本
    sample_text = """
    人工智能是计算机科学的一个分支，它企图了解智能的实质，
    并生产出一种新的能以人类智能相似的方式做出反应的智能机器。
    人工智能的研究包括机器学习、计算机视觉、自然语言处理等领域。
    硅基流动是一家专注于AI技术研究和应用的公司，提供高质量的AI服务。
    """
    
    try:
        # 处理并存储文本
        vector_store = search_engine.process_and_store(sample_text)
        
        # 搜索示例
        query = "硅基流动"
        results = search_engine.search_similar(query)
        
        # 打印结果
        for doc, score in results:
            print(f"相似度分数: {score}")
            print(f"内容: {doc.page_content}\n")
    except Exception as e:
        print(f"程序执行过程中发生错误: {str(e)}")

if __name__ == "__main__":
    main()






